#!/usr/bin/env python
# @Author  : Bright (brt2@qq.com)
# @Create  : 2025-09-05
# @Update  : 2025-09-18 / 0.2.3

import os
import time
import base64
import re
import openai

from .putils import make_logger

class ChatModelApi:
    """ 用于兼容本地vLLM & ollama的API接口 """
    USE_OLD_API = True

    # ---------- 初始化 ----------
    def __init__(self, logger_level=20, logger_path=None):
        self.logger = make_logger("ChatModelApi", logger_path, level=logger_level)
        self._init_prompt__()

    def connect(self, api_base=None, api_key=None, model_tag=None, **kwargs):
        base_url = api_base or os.environ.get("OPENAI_API_URL")
        if not base_url:
            base_url = "https://api.openai.com/v1"
        api_key = api_key or os.environ.get("OPENAI_API_KEY", "")
        model_name = model_tag or os.environ.get("OPENAI_MODEL_NAME")
        self.logger.info(f"Connect to【{base_url}】with model【{model_name}】")

        self.client = openai.OpenAI(base_url=base_url, api_key=api_key)

        self.params = {
            "model": model_name,
            "stream": False,
        }
        self.option = {
            "temperature": 0.15,
            "top_p": 0.9,
            **kwargs
        }

    def _init_prompt__(self):
        self.system_input = "You are a helpful assistant."
        self.user_input = ""
        self.assistant_input = ""
        self.message_tmpl = [
            {"role": "system", "content": self.system_input},
            # {"role": "user", "images": [self.load_img("sofa.png")],
            #          "content": self.user_input},
            # {"role": "assistant", "content": self.assistant_input},
        ]
        self.reset_chat_history()

    def reset_chat_history(self):
        self.message_history = self.message_tmpl.copy()

    # ---------- 工具函数 ----------
    @staticmethod
    def load_img(image_path: str) -> str:
        """加载图片并转为 base64"""
        with open(image_path, 'rb') as image_file:
            binary_data = image_file.read()
        return base64.b64encode(binary_data).decode('utf-8')

    def _build_input(self, question=None, img_data: list = None):
        """生成 input 内容（文本 + 图片）"""
        content = []
        if question:
            content.append({"type": "text", "text": question})
        if img_data:
            for img in img_data:
                content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}})
        return {"role": "user", "content": content}

    # ---------- Generate ----------
    def generate(self, question=None, img_data: list = None, ret_str=True):
        """一次性生成，不记忆上下文"""
        message_list = self.message_tmpl.copy()
        message_list.append(self._build_input(question, img_data))
        self.logger.debug(f">>> Prompt: {message_list}")

        if self.USE_OLD_API:
            response = self.client.chat.completions.create(
                messages=message_list,
                **self.params,
                **self.option
            )
        else:
            response = self.client.responses.create(
                input=message_list,
                **self.params,
                **self.option
            )

        return self.parse_generate(response) if ret_str else response

    def parse_generate(self, response):
        """解析 generate 的结果"""
        self.logger.debug(f">>> Response: {response}")
        if self.USE_OLD_API:
            # return response.choices[0].text  # completions
            return response.choices[0].message.content
        else:
            return response.output[0].content[0].text

    def __generate_until_finish_oldapi(self, question=None, img_data: list = None, ret_str=True, wait_time=5, max_try=10):
        """ 由于API可能频率受到限制，持续catch异常，直至成功完成调用 """
        for i in range(max_try):
            try:
                return self.generate(question, img_data, ret_str)
            except openai.error.ServiceUnavailableError:
                self.logger.warning('OpenAI ServiceUnavailableError.\tWill try again.')
            except openai.error.RateLimitError:
                self.logger.warning('OpenAI RateLimitError.\tWill try again.')
            except openai.error.APIError as e:
                self.logger.warning('OpenAI Bad Gateway Error.\tWill try again.')
            except openai.error.InvalidRequestError as e:
                # self.logger.debug(e)
                self.logger.warning('Input too long. Will shrink the prompting examples.')
                self.option['top_k'] -= 1
                continue
            except Exception as e:
                self.logger.error(f"Unknown Error: {e}")
            finally:
                time.sleep(wait_time)

        raise Exception(f"Failed to complete chat after max_try={max_try} times.")

    def generate_until_finish(self, question=None, img_data: list = None, ret_str=True, wait_time=5, max_try=10):
        """ 由于API可能频率受到限制，持续catch异常，直至成功完成调用 """
        for i in range(max_try):
            try:
                return self.generate(question, img_data, ret_str)
            except openai.OpenAIError as e:
                # 处理OpenAI SDK v1.x版本的异常
                error_type = e.__class__.__name__
                if error_type in ['ServiceUnavailableError', 'APIConnectionError']:
                    self.logger.warning(f'OpenAI {error_type}.\tWill try again.')
                elif error_type == 'RateLimitError':
                    self.logger.warning('OpenAI RateLimitError.\tWill try again.')
                elif error_type == 'BadRequestError':
                    # InvalidRequestError在v1中被重命名为BadRequestError
                    self.logger.warning('Input too long or invalid. Will shrink the prompting examples.')
                    if 'top_k' in self.option:
                        self.option['top_k'] -= 1
                    continue
                else:
                    self.logger.warning(f'OpenAI {error_type}: {str(e)}\tWill try again.')
            except Exception as e:
                self.logger.error(f"Unknown Error: {e}")
            finally:
                time.sleep(wait_time)

        raise Exception(f"Failed to complete chat after max_try={max_try} times.")

    # ---------- Chat ----------
    def chat(self, question=None, img_data: list = None, ret_str=True):
        """多轮对话，自动缓存上下文"""
        dict_role_with_content = self._build_input(question, img_data)
        self.message_history.append(dict_role_with_content)
        self.logger.debug(f">>> Prompt: {self.message_history}")

        if self.USE_OLD_API:
            response = self.client.chat.completions.create(
                messages=self.message_history,
                **self.params,
                **self.option,
            )
            msg = self.parse_chat(response)
        else:
            response = self.client.completions.create(
                input=self.message_history,
                **self.params,
                **self.option,
            )
            msg = self.parse_generate(response)

        # 将模型回复加入上下文
        self.message_history.append({"role": "assistant", "content": msg})
        return msg if ret_str else response

    def parse_chat(self, response):
        """ deprecated: 用于解析client.chat.completions的结果（老版本接口） """
        self.logger.debug(f">>> Response: {response}")
        return response.choices[0].message.content

    def parse_as_json(self, json_markdown_str):
        """ 从Markdown字符串中提取JSON内容 """
        self.logger.debug(f">>> Input: {json_markdown_str}")
        pattern = r'```json\s*([^`]+)\s*```' # Match the json object
        match = re.search(pattern, json_markdown_str, re.DOTALL).group(1)
        # json_obj_new = json.loads(match)
        if match is None:
            self.logger.warning(f"Failed to parse JSON from LLM response")
        return match

    # ---------- 图像生成 ----------
    def generate_image(self, prompt, size="1024x1024", n=1, quality="standard", style="vivid"):
        """使用OpenAI API生成图像"""
        if not self.client:
            raise ValueError("OpenAI client not initialized. Please call connect() first.")

        # 构建请求参数
        payload = {
            "prompt": prompt,
            "n": n,
            "size": size,
            "quality": quality,
            "style": style
        }

        try:
            self.logger.debug(f">>> Generating image with prompt: {prompt}")
            response = self.client.images.generate(**payload)
            
            # 解析响应
            image_data = []
            for data in response.data:
                image_info = {
                    "url": getattr(data, 'url', None),
                    "b64_json": getattr(data, 'b64_json', None),
                    "revised_prompt": getattr(data, 'revised_prompt', None)
                }
                image_data.append(image_info)
            
            result = {
                "prompt": prompt,
                "model": self.params.get("model", "dall-e-3"),
                "size": size,
                "n": n,
                "quality": quality,
                "style": style,
                "images": image_data
            }
            
            self.logger.info("Image generation successful")
            return result
            
        except Exception as e:
            error_msg = f"Image generation failed: {str(e)}"
            self.logger.error(error_msg)
            raise Exception(error_msg)


if __name__ == "__main__":
    api = ChatModelApi()
    api.connect()

    response = api.generate("做个简短的自我介绍")
    print("Generate 输出:\n", response, "\n")

    # --------- 图像理解 ---------
    if False:
        path_img = "tmp/anno_images/4.jpg"
        img_b64 = api.load_img(path_img)

        response = api.generate("用中文详细描述图片中红色包围框标注的物体", [img_b64])
        print("Generate 输出:\n", response, "\n" * 2)

    # --------- chat 模式 ---------
    response = api.chat("简单说明什么是cosplay")
    print("Chat 输出:\n", response, "\n")

    # --------- chat again（上下文记忆）---------
    response = api.chat("请把上面的结果翻译成英文")
    print("Chat Again 输出:\n", response)
